import open3d as o3d
import os 
import numpy as np

WeldPoint = np.zeros((257, 5, 1024, 3))
WeldPoint1 = np.zeros((257, 5, 512, 3))
WeldPoint2 = np.zeros((257, 5, 512, 3))

WeldingCoordinates = np.array([
[3.94, -8.55, 2],
[0.14, -8.55, 2],
[-3.85, -8.55, 2],
[1.44, 8.44, 2],
[-0.75, 8.44, 2],
]) 

WeldingCoordinates_New = np.array([
[4.17, -8.55, 2],    
[0.17, -8.55, 2],
[-3.62, -8.55, 2],
[0.87, 8.55, 2],
[-1.13, 8.55, 2],
]) 

def display_inlier_outlier(cloud, ind):
    inlier_cloud = cloud.select_by_index(ind)
    outlier_cloud = cloud.select_by_index(ind, invert=True)  # invert利用索引反选

    print("Showing outliers (red) and inliers (gray): ")
    outlier_cloud.paint_uniform_color([1, 0, 0])
    inlier_cloud.paint_uniform_color([0.8, 0.8, 0.8])
    o3d.visualization.draw_geometries([inlier_cloud, outlier_cloud])  # 可视化降采样的点、未被选中的点

def cropWelding(padCloud, WeldingCoordinates=WeldingCoordinates):
    arrpad = np.asarray(padCloud.points)
    weldpoint = [0]*5 
    for i in range(5):
        if(i == 3):
            xfilter = arrpad[arrpad[:, 0] > -2.15]
            xfilter = xfilter[xfilter[:, 0] < 0.33]
        elif(i == 4):
            xfilter = arrpad[arrpad[:, 0] > 0.33]
            xfilter = xfilter[xfilter[:, 0] < 2.16]
        else:
            xfilter = arrpad[arrpad[:, 0] > (WeldingCoordinates[i, 0] - 1.1)]
            xfilter = xfilter[xfilter[:, 0] < (WeldingCoordinates[i, 0] + 1.1)]
        yfiler =  xfilter[xfilter[:, 1] > (WeldingCoordinates[i, 1]-1.1)]
        yfiler =  yfiler[yfiler[:, 1] < (WeldingCoordinates[i, 1]+1.1)]
        zfiler = yfiler[yfiler[:, 2] < WeldingCoordinates[i, 2]]
        weldpoint[i] = zfiler

    weld = [0]*5
    for i in range(5):
        weld[i] = o3d.geometry.PointCloud()
        weld[i].points =  o3d.utility.Vector3dVector(weldpoint[i])
    return weld

def farthest_point_sample(xyz, npoint):
    """
    Input:
        xyz: pointcloud data, [B, N, 3]
        npoint: number of samples
    Return:
        centroids: sampled pointcloud index, [B, npoint]
    """
    B, N, C = xyz.shape
    centroids = np.zeros((B, npoint), dtype=np.long)
    distance = np.ones((B, N)) * 1e10
    farthest = np.random.randint(0, N, (B,), dtype=np.long)# 随机
    batch_indices = np.arange(B, dtype=np.long)# 随机索引
    for i in range(npoint):
        centroids[:, i] = farthest
        centroid = xyz[batch_indices, farthest, :].reshape((B, 1, 3))
        dist = np.sum((xyz - centroid) ** 2, -1)
        mask = dist < distance
        distance[mask] = dist[mask]
        farthest = np.argmax(distance, -1)
    return centroids

# data_path = 'D:/mycode/dataset/20201110/' 
npoint = 1024
# 读取点云
# sampledirs = os.listdir(data_path)
#样本二
tar = o3d.io.read_point_cloud('SampleTarget.pcd')
tar1 = o3d.io.read_point_cloud('SampleTarget2.pcd')
taridx = [27, 35, 55, 63, 110, 112, 117, 119, 144, 150, 151, 153, 159, 164, 169, 174, 179, 180,\
    191, 193, 203, 207, 208, 213, 215, 217, 221, 222, 225, 226, 227, 233, 237, 243, 245, 246, 247, 248]
for n in range(257):
    source1 = o3d.io.read_point_cloud('D:\WorkSpace\python\PointPreproces\data\points_200_'+str(n)+'.pcd')
    source2 = o3d.io.read_point_cloud('D:\WorkSpace\python\PointPreproces\data\points_201_'+str(n)+'.pcd')

    source1_np = source1.points
    source2_np = source2.points
        
    sampleweld1s = cropWelding(source1, WeldingCoordinates)
    sampleweld2s = cropWelding(source2, WeldingCoordinates)
    # o3d.io.write_point_cloud('data/icp_cen/points_200_'+str(n)+'.pcd', source1)
    # o3d.io.write_point_cloud('data/icp_cen/points_201_'+str(n)+'.pcd', source2)
    # o3d.visualization.draw_geometries(sampleweld1s[:3])

    nidx = 0
    for weld1, weld2 in zip(sampleweld1s, sampleweld2s):
        weld = weld1 + weld2
        # 去重
        # np_weld_points = np.asarray(weld.points)
        # weld.points = o3d.utility.Vector3dVector(np_weld_points - np.mean(np_weld_points, axis=0))
        # diameter = np.linalg.norm(np.asarray(weld.get_max_bound()) - np.asarray(weld.get_min_bound()))
        # camera = [0, 0, diameter]
        # radius = diameter * 300
        # _, pt_map = weld.hidden_point_removal(camera, radius)
        # weld = weld.select_by_index(pt_map)
        # arrWeld1, arrWeld2 = np.asarray(weld1.points), np.asarray(weld2.points)
        # arrWeld = np.concatenate((arrWeld1,arrWeld2),axis=0)
        # o3d.visualization.draw_geometries([weld])
        arrWeld = np.asarray(weld.points)

        xyz = np.expand_dims(arrWeld, axis=0)
        try:
            idx = farthest_point_sample(xyz, npoint)
        except Exception as e:
            print('not enough point')
            exit()

        pcd = [0]*len(xyz)
        for i in range(len(xyz)):
            pcd[i] = o3d.geometry.PointCloud()
            pcd[i].points = o3d.utility.Vector3dVector(xyz[i, idx[i]])
        # o3d.visualization.draw_geometries(pcd)
        WeldPoint[n, nidx] = pcd[0].points
    
        nidx += 1
        
    np.save('all_point_new3.npy', WeldPoint)
    np.save('all_point_flatten_new3.npy', WeldPoint.reshape((-1,npoint,3)))
#     for weld1, weld2 in zip(sampleweld1s, sampleweld2s):
#         np_weld_points1, np_weld_points2 = np.asarray(weld1.points), np.asarray(weld2.points)
#         np_weld_points1 = np_weld_points1 - np.mean(np_weld_points1, axis=0)
#         np_weld_points2 = np_weld_points2 - np.mean(np_weld_points2, axis=0)

#         xyz1 = np.expand_dims(np_weld_points1, axis=0)
#         try:
#             idx1 = farthest_point_sample(xyz1, 512)
#         except Exception as e:
#             print('not enough point')
#             exit()
#         xyz2 = np.expand_dims(np_weld_points2, axis=0)
#         try:
#             idx2 = farthest_point_sample(xyz2, 512)
#         except Exception as e:
#             print('not enough point')
#             exit()

#         WeldPoint1[n, nidx] = xyz1[0, idx1[0]]
#         WeldPoint2[n, nidx] = xyz2[0, idx2[0]]
#         nidx += 1
    
    print('sample %d is ok' % (n+1))
# np.save('all_point_flatten1.npy', WeldPoint1.reshape((-1,512,3)))
# np.save('all_point_flatten2.npy', WeldPoint2.reshape((-1,512,3)))
print(':: all sample have been saving')
# o3d.io.write_point_cloud('D:/mycode/data_process/res.pcd',pcd)